import argparse

import torch
import yaml
#from torchstat import stat


from ResNet import ResNet
from ResNet_2 import ResNet_2



from dataset import initialize_dataset
from train_test import Training
from trainAndTestWithSAM import TrainingWithSAM

"""Device Selection"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

""" Initialize model based on command line argument """
model_parser = argparse.ArgumentParser(description='Image Classification Using PyTorch', usage='[option] model_name')
model_parser.add_argument('--model', type=str, required=True)
model_parser.add_argument('--model_save', type=bool, required=False)
model_parser.add_argument('--checkpoint', type=bool, required=False)
model_parser.add_argument('--sam', type=bool, required=False)
args = model_parser.parse_args()

"""Loading Config File"""
try:
    stream = open("config.yaml", 'r')
    config = yaml.safe_load(stream)
except FileNotFoundError:
    print("Config file missing")

"""Dataset Initialization"""
data_initialization = initialize_dataset(image_resolution=config['parameters']['image_resolution'], batch_size=config['parameters']['batch_size'], 
                      )
train_dataloader, test_dataloader = data_initialization.load_dataset(transform=True)

input_channel = next(iter(train_dataloader))[0].shape[1]
#n_classes = len(torch.unique(next(iter(train_dataloader))[1]))
n_classes = config['parameters']['n_classes']

"""Model Initialization"""



if args.model == 'resnet':
    model = ResNet(input_channel=input_channel, n_classes=n_classes).to(device)
elif args.model == 'resnet2':
    model = ResNet_2(input_channel=input_channel, n_classes=n_classes).to(device) 





#print(device)

print(f'Total Number of Parameters of {args.model.capitalize()} is {round((sum(p.numel() for p in model.parameters()))/1000000, 2)}M')
if not args.sam:
    trainer = Training(model=model, optimizer=config['parameters']['optimizer'], learning_rate=config['parameters']['learning_rate'], 
                train_dataloader=train_dataloader, num_epochs=config['parameters']['num_epochs'],test_dataloader=test_dataloader,
                model_name=args.model, model_save=args.model_save, checkpoint=args.checkpoint)
    trainer.runner()
else:
    trainer = TrainingWithSAM(model=model, optimizer=config['parameters']['optimizer'], learning_rate=config['parameters']['learning_rate'], 
                train_dataloader=train_dataloader, num_epochs=config['parameters']['num_epochs'],test_dataloader=test_dataloader,
                model_name=args.model, model_save=args.model_save, checkpoint=args.checkpoint)
    trainer.runner()
    
# Calculate FLops and Memory Usage.
# model.to('cpu')
# dummy_input = (input_channel, config['parameters']["image_resolution"], config['parameters']["image_resolution"])
# print(stat(model, dummy_input))
